# encoding=utf-8

import os
import cv2
import numpy as np
from PIL import Image
from seg_cuda.CudaMatrixTools import MatBuilder
from seg_common.logging import ConsoleService

from seg_system.vascular.service.VascularToolsForBatch.VascularEachProcessor.VascularForBatchBase import \
    VascularForBatchBase


class CTBDForBatch(VascularForBatchBase):
    def CTBD_base_parser(self, **kwargs):
        # 不好一次性导入，所以采用这种方案
        file_name_list = kwargs.get("file_name_list")
        refine_save_path = kwargs.get("refine_save_path")
        seg_save_path = kwargs.get("seg_save_path")

        if (file_name_list is None) or (refine_save_path is None) or (seg_save_path is None):
            raise Exception("Input name_list or save path is Empty, please check user is exits, "
                            "or forget input to algorithm")

        return file_name_list, refine_save_path, seg_save_path

    def process_with_python(self, big_matrix: np.ndarray, usable_matrix: np.ndarray, each_matrix_shape: tuple,
                            stride: int, **kwargs):
        # 尽量保证和VascularTools/CTBD部分一致

        # CTBD的python部分，不能进行整体运算, one_by_one只能是True
        one_by_one = kwargs.get("one_by_one", True)
        if not one_by_one:
            ConsoleService.console_log("CTBD in python is not support cal in once", ConsoleService.WARNING)

        # 不好一次性导入，所以采用这种方案
        file_name_list, refine_save_path, seg_save_path = self.CTBD_base_parser(**kwargs)

        usable_shape = usable_matrix.shape
        gray_big_matrix = cv2.cvtColor(big_matrix, cv2.COLOR_BGR2GRAY)
        CTBD_list = []
        for i in range(usable_shape[0]):
            for j in range(usable_shape[1]):
                if usable_matrix[i][j] == 0:
                    continue

                h_start = i * (each_matrix_shape[0] + stride)
                w_start = j * (each_matrix_shape[1] + stride)

                # each_mat = img
                img1 = gray_big_matrix[h_start: h_start + each_matrix_shape[0],
                       w_start: w_start + each_matrix_shape[1]]
                num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(img1, connectivity=8)

                file_name_index = i * usable_shape[1] + j
                file_name = file_name_list[file_name_index]

                img11 = cv2.imread(os.path.join(refine_save_path, file_name), 0)
                img111 = self.getSkeletonIntersection(img11)
                imgsu = cv2.imread(os.path.join(seg_save_path, file_name))
                imgss = cv2.cvtColor(imgsu, cv2.COLOR_BGR2GRAY)

                num__labels, labels, stats, centroids = cv2.connectedComponentsWithStats(imgss, connectivity=8)
                NUMBER = num__labels - num_labels - len(img111)
                CTBD_list.append(NUMBER)

        return CTBD_list

    def process_with_cxx(self, big_matrix: np.ndarray, usable_matrix: np.ndarray, each_matrix_shape: tuple, stride: int,
                         **kwargs):
        # CTBD的C++部分，不能进行整体运算, one_by_one只能是True
        # 等到后面把cuda版本改好后，才能进行整体运算
        one_by_one = kwargs.get("one_by_one", True)
        if not one_by_one:
            ConsoleService.console_log("CTBD in CXX is not support cal in once, please wait for CTBD cuda update",
                                       ConsoleService.WARNING)

        file_name_list, refine_save_path, seg_save_path = self.CTBD_base_parser(**kwargs)
        CTBD_list = self.cxxPyVascular.processCTBD(big_matrix, usable_matrix.tolist(), each_matrix_shape[0],
                                                   each_matrix_shape[1],
                                                   stride, file_name_list, refine_save_path, seg_save_path)

        return CTBD_list

    def batch_save(self, file_name: list, save_path: str, process_output, **kwargs):
        return process_output
